import gdal
import os

def raster_copy_with_nodata(s_fh, s_xoff, s_yoff, s_xsize, s_ysize, s_band_n,
                            t_fh, t_xoff, t_yoff, t_xsize, t_ysize, t_band_n,
                            nodata):
    """Copy a band of raster into the output file with nodata values.

       Function copied from gdal_merge.py
    """
    try:
        import numpy as Numeric
    except ImportError:
        import Numeric

    s_band = s_fh.GetRasterBand(s_band_n)
    t_band = t_fh.GetRasterBand(t_band_n)
    data_src = s_band.ReadAsArray(s_xoff, s_yoff, s_xsize, s_ysize,
                                  t_xsize, t_ysize)
    # data_src = min_max_normalization_fill(data_src, 2, nodata)
    data_dst = t_band.ReadAsArray(t_xoff, t_yoff, t_xsize, t_ysize)
    nodata_test = Numeric.equal(data_src, nodata)
    to_write = Numeric.choose(nodata_test, (data_src, data_dst))
    t_band.WriteArray(to_write, t_xoff, t_yoff)

    return 0


def raster_copy(s_fh, s_xoff, s_yoff, s_xsize, s_ysize, s_band_n,
                t_fh, t_xoff, t_yoff, t_xsize, t_ysize, t_band_n,
                nodata=None):
    """Copy a band of raster into the output file.

       Function copied from gdal_merge.py
    """
    if nodata is not None:
        return raster_copy_with_nodata(s_fh, s_xoff, s_yoff, s_xsize, s_ysize,
                                       s_band_n, t_fh, t_xoff, t_yoff, t_xsize,
                                       t_ysize, t_band_n, nodata)
    s_band = s_fh.GetRasterBand(s_band_n)
    t_band = t_fh.GetRasterBand(t_band_n)
    data = s_band.ReadRaster(s_xoff, s_yoff, s_xsize, s_ysize,
                             t_xsize, t_ysize, t_band.DataType)
    t_band.WriteRaster(t_xoff, t_yoff, t_xsize, t_ysize, data, t_xsize,
                       t_ysize, t_band.DataType)
    return 0


def copy_into(t_fh, in_ds, s_band=1, t_band=1, nodata_arg=None):
    """Copy this files image into target file.
    This method will compute the overlap area of the file_info objects
    file, and the target gdal.Dataset object, and copy the image data
    for the common window area.  It is assumed that the files are in
    a compatible projection. no checking or warping is done.  However,
    if the destination file is a different resolution, or different
    image pixel type, the appropriate resampling and conversions will
    be done (using normal GDAL promotion/demotion rules).

    :param t_fh: gdal.Dataset object for the file into which some or all
                 of this file may be copied.
    :param s_band:
    :param t_band:
    :param nodata_arg:

    :return: 1 on success (or if nothing needs to be copied), and zero one
             failure.

    """
    geotransform = in_ds.GetGeoTransform()
    xsize = in_ds.RasterXSize
    ysize = in_ds.RasterYSize
    ulx = geotransform[0]
    uly = geotransform[3]
    lrx = ulx + geotransform[1] * xsize
    lry = uly + geotransform[5] * ysize
    t_geotransform = t_fh.GetGeoTransform()
    t_ulx = t_geotransform[0]
    t_uly = t_geotransform[3]
    t_lrx = t_geotransform[0] + t_fh.RasterXSize * t_geotransform[1]
    t_lry = t_geotransform[3] + t_fh.RasterYSize * t_geotransform[5]

    # figure out intersection region
    tgw_ulx = max(t_ulx, ulx)
    tgw_lrx = min(t_lrx, lrx)
    if t_geotransform[5] < 0:
        tgw_uly = min(t_uly, uly)
        tgw_lry = max(t_lry, lry)
    else:
        tgw_uly = max(t_uly, uly)
        tgw_lry = min(t_lry, lry)

    # do they even intersect?
    if tgw_ulx >= tgw_lrx:
        return 1
    if t_geotransform[5] < 0 and tgw_uly <= tgw_lry:
        return 1
    if t_geotransform[5] > 0 and tgw_uly >= tgw_lry:
        return 1

    # compute target window in pixel coordinates.
    tw_xoff = int((tgw_ulx - t_geotransform[0]) / t_geotransform[1] + 0.1)
    tw_yoff = int((tgw_uly - t_geotransform[3]) / t_geotransform[5] + 0.1)
    tw_xsize = int((tgw_lrx - t_geotransform[0]) / t_geotransform[1] + 0.5) - tw_xoff
    tw_ysize = int((tgw_lry - t_geotransform[3]) / t_geotransform[5] + 0.5) - tw_yoff

    if tw_xsize < 1 or tw_ysize < 1:
        return 1

    # Compute source window in pixel coordinates.
    sw_xoff = int((tgw_ulx - geotransform[0]) / geotransform[1])
    sw_yoff = int((tgw_uly - geotransform[3]) / geotransform[5])
    sw_xsize = int((tgw_lrx - geotransform[0])
                   / geotransform[1] + 0.5) - sw_xoff
    sw_ysize = int((tgw_lry - geotransform[3])
                   / geotransform[5] + 0.5) - sw_yoff

    if sw_xsize < 1 or sw_ysize < 1:
        return 1
    return \
        raster_copy(in_ds, sw_xoff, sw_yoff, sw_xsize, sw_ysize, s_band,
                    t_fh, tw_xoff, tw_yoff, tw_xsize, tw_ysize, t_band,
                    nodata_arg)


def copy_into_round_up(t_fh, in_ds, s_band=1, t_band=1, nodata_arg=None):
    """Copy this files image into target file.
    This method will compute the overlap area of the file_info objects
    file, and the target gdal.Dataset object, and copy the image data
    for the common window area.  It is assumed that the files are in
    a compatible projection. no checking or warping is done.  However,
    if the destination file is a different resolution, or different
    image pixel type, the appropriate resampling and conversions will
    be done (using normal GDAL promotion/demotion rules).

    :param t_fh: gdal.Dataset object for the file into which some or all
                 of this file may be copied.
    :param s_band:
    :param t_band:
    :param nodata_arg:

    :return: 1 on success (or if nothing needs to be copied), and zero one
             failure.

    """
    geotransform = in_ds.GetGeoTransform()
    xsize = in_ds.RasterXSize
    ysize = in_ds.RasterYSize

    geotransform0 = round_up(geotransform[0], 3)
    geotransform3 = round_up(geotransform[3], 3)

    ulx = geotransform0
    uly = geotransform3
    lrx = round_up(ulx + geotransform[1] * xsize, 3)
    lry = round_up(uly + geotransform[5] * ysize, 3)
    t_geotransform = t_fh.GetGeoTransform()

    t_geotransform0 = round_up(t_geotransform[0], 3)
    t_geotransform3 = round_up(t_geotransform[3], 3)

    t_ulx = t_geotransform0
    t_uly = t_geotransform3
    t_lrx = t_geotransform0 + t_fh.RasterXSize * t_geotransform[1]
    t_lry = t_geotransform3 + t_fh.RasterYSize * t_geotransform[5]

    # figure out intersection region
    tgw_ulx = max(t_ulx, ulx)
    tgw_lrx = min(t_lrx, lrx)
    if t_geotransform[5] < 0:
        tgw_uly = min(t_uly, uly)
        tgw_lry = max(t_lry, lry)
    else:
        tgw_uly = max(t_uly, uly)
        tgw_lry = min(t_lry, lry)

    # do they even intersect?
    if tgw_ulx >= tgw_lrx:
        return 1
    if t_geotransform[5] < 0 and tgw_uly <= tgw_lry:
        return 1
    if t_geotransform[5] > 0 and tgw_uly >= tgw_lry:
        return 1

    # compute target window in pixel coordinates.
    tw_xoff = int((tgw_ulx - t_geotransform0) / t_geotransform[1] + 0.1)
    tw_yoff = int((tgw_uly - t_geotransform3) / t_geotransform[5] + 0.1)
    tw_xsize = int((tgw_lrx - t_geotransform0) / t_geotransform[1] + 0.5) - tw_xoff
    tw_ysize = int((tgw_lry - t_geotransform3) / t_geotransform[5] + 0.5) - tw_yoff

    if tw_xsize < 1 or tw_ysize < 1:
        return 1

    # Compute source window in pixel coordinates.
    sw_xoff = int((tgw_ulx - geotransform0) / t_geotransform[1])
    sw_yoff = int((tgw_uly - geotransform3) / t_geotransform[5])
    sw_xsize = int((tgw_lrx - geotransform0)
                   / t_geotransform[1] + 0.5) - sw_xoff
    sw_ysize = int((tgw_lry - geotransform[3])
                   / t_geotransform[5] + 0.5) - sw_yoff

    if sw_xsize < 1 or sw_ysize < 1:
        return 1
    return \
        raster_copy(in_ds, sw_xoff, sw_yoff, sw_xsize, sw_ysize, s_band,
                    t_fh, tw_xoff, tw_yoff, tw_xsize, tw_ysize, t_band,
                    nodata_arg)


def get_extent(in_fn):
    ds = gdal.Open(in_fn)
    geotrans = list(ds.GetGeoTransform())
    xsize = ds.RasterXSize
    ysize = ds.RasterYSize
    min_x = geotrans[0]
    max_y = geotrans[3]
    max_x = geotrans[0] + xsize * geotrans[1]
    min_y = geotrans[3] + ysize * geotrans[5]
    ds = None
    return min_x, max_y, max_x, min_y


def find_layers_files(dir_name, format):
    file_list = []
    for home, dirs, files in os.walk(dir_name):
        for filename in sorted(files):
            if filename.endswith(format):
                file_list.append(os.path.join(home, filename))
    return file_list


def merge_modis_to_tiff(tiff_list, output_dir, fill):
    min_x, max_y, max_x, min_y = get_extent(tiff_list[0])
    for in_fn in tiff_list[1:]:
        minx, maxy, maxx, miny = get_extent(in_fn)
        min_x = round_up(min(min_x, minx), 3)
        min_y = round_up(min(min_y, miny), 3)
        max_x = round_up(max(max_x, maxx), 3)
        max_y = round_up(max(max_y, maxy), 3)
    in_ds = gdal.Open(tiff_list[0])
    geotrans = list(in_ds.GetGeoTransform())

    geotrans = [min_x, geotrans[1], 0, max_y, 0, geotrans[5]]
    xsize = int((max_x - min_x) / geotrans[1] + 0.5)
    ysize = int((min_y - max_y) / geotrans[5] + 0.5)
    # band_type = in_ds.GetRasterBand(1).DataType
    driver = gdal.GetDriverByName('GTiff')
    out_ds = driver.Create(output_dir, xsize, ysize, 1, gdal.GDT_Float32)
    # fill = float(in_ds.GetMetadata()['_FillValue'])
    out_ds.SetGeoTransform(geotrans)
    out_ds.SetProjection(in_ds.GetProjection())
    out_ds.GetRasterBand(1).SetNoDataValue(fill)
    out_ds.GetRasterBand(1).Fill(fill)
    for in_fns in tiff_list:
        in_ds = gdal.Open(in_fns)

        print(in_fns)

        copy_into(out_ds, in_ds, 1, 1, fill)


def round_up(number,power=0):
    """
    实现精确四舍五入，包含正、负小数多种场景
    :param number: 需要四舍五入的小数
    :param power: 四舍五入位数，支持0-∞
    :return: 返回四舍五入后的结果
    """
    digit = 10 ** power
    num2 = float(int(number * digit))
    # 处理正数，power不为0的情况
    if number>=0 and power !=0:
        tag = number * digit - num2 + 1 / (digit * 10)
        if tag>=0.5:
            return (num2+1)/digit
        else:
            return num2/digit
    # 处理正数，power为0取整的情况
    elif  number>=0 and power==0 :
        tag = number * digit - int(number)
        if tag >= 0.5:
            return (num2 + 1) / digit
        else:
            return num2 / digit
    # 处理负数，power为0取整的情况
    elif power==0 and number<0:
        tag = number * digit - int(number)
        if tag <= -0.5:
            return (num2 - 1) / digit
        else:
            return num2 / digit
    # 处理负数，power不为0的情况
    else:
        tag = number * digit - num2 - 1 / (digit * 10)
        if tag <= -0.5:
            return (num2-1)/digit
        else:
            return num2/digit